#include "DXUT.h"
#include <PrecompiledHeader.h>
#include "Include.h"
#include "Shader.h"
using namespace d3d11;

//--------------------------------------------------------------------------------------
// Use this until D3DX11 comes online and we get some compilation helpers - DXUT
//--------------------------------------------------------------------------------------
HRESULT CompileShaderFromFile( LPCWSTR szFileName, LPCSTR szEntryPoint, LPCSTR szShaderModel, ID3DBlob** ppBlobOut )
{
    HRESULT hr = S_OK;
	TRACEW(L"CompileShaderFromFile( \'%ls\' ", szFileName);
	TRACE( "\'%s\', \'%s\', ... )\n", szEntryPoint, szShaderModel);
    // find the file
    //WCHAR str[MAX_PATH];
    //DXUTFindDXSDKMediaFileCch( str, MAX_PATH, szFileName ) ;

    // open the file
    HANDLE hFile = CreateFile( szFileName, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING,
                               FILE_FLAG_SEQUENTIAL_SCAN, NULL );
    if( INVALID_HANDLE_VALUE == hFile )
        return E_FAIL;

    // Get the file size
    LARGE_INTEGER FileSize;
    GetFileSizeEx( hFile, &FileSize );

    // create enough space for the file data
    BYTE* pFileData = new BYTE[ FileSize.LowPart ];
    if( !pFileData )
        return E_OUTOFMEMORY;

    // read the data in
    DWORD BytesRead;
    if( !ReadFile( hFile, pFileData, FileSize.LowPart, &BytesRead, NULL ) )
        return E_FAIL; 

    CloseHandle( hFile );

    // Compile the shader
    ID3DBlob* pErrorBlob;
	Include	includedInterface;
    hr = D3DCompile( pFileData, FileSize.LowPart, "none", NULL, &includedInterface, szEntryPoint, szShaderModel, D3D10_SHADER_ENABLE_STRICTNESS, 0, ppBlobOut, &pErrorBlob );

    delete []pFileData;

    if( FAILED(hr) )
    {
		TRACEW(L"D3DCompile Failed in CompileShaderFromFile\n");
        //OutputDebugStringA( (char *)pErrorBlob->GetBufferPointer() );
		TRACE("[SHADER COMPILE ERROR] %s\n", (char *)pErrorBlob->GetBufferPointer());
        SAFE_RELEASE( pErrorBlob );
        return hr;
    }
    SAFE_RELEASE( pErrorBlob );

    return S_OK;
}


Shader::Shader()
:m_pBlob(NULL), m_Type(UNKNOWN)
{	
	m_pD3DDevice9 = DXUTGetD3D9Device();
	m_pD3DDevice11 = DXUTGetD3D11Device();
	m_pBlob = NULL;
	m_Shader.m_pVertexShader11 = NULL;
	m_Shader.m_pPixelShader11 = NULL;
	m_Shader.m_pVertexShader9 = NULL;
	m_Shader.m_pPixelShader9 = NULL;
}

Shader::~Shader()
{
	switch(m_Type)
	{
	case VERTEX11:
		SAFE_RELEASE(m_pBlob);
		SAFE_RELEASE(m_Shader.m_pVertexShader11);
		break;
	case PIXEL11:
		SAFE_RELEASE(m_pBlob);
		SAFE_RELEASE(m_Shader.m_pPixelShader11);
		break;
	case VERTEX9:
		SAFE_RELEASE(m_Shader.m_pVertexShader9);
		break;
	case PIXEL9:
		SAFE_RELEASE(m_Shader.m_pPixelShader9);
		break;
	};
}

// TODO: change to correct vs_4_0 / vs_5_0 etc level based on machine capability
void Shader::ConstructShader(Type type, LPCWSTR filepath, LPCSTR entryFunction)
{
	CASSERT(m_Shader.m_pVertexShader11 == NULL, "m_pShader is already set.");
	CASSERT(m_pBlob == NULL, "m_pBlob is already set.");

	m_Type = type;

	switch(type)
	{
	case VERTEX11:
		{
			HRESULT hr = S_OK;
			/*
			if (FAILED(hr = D3DX10CompileFromFile()))
			{
				CASSERT(m_pBlob != NULL, "CompileShaderFromFile failed.");
			}*/
			if(FAILED(hr = CompileShaderFromFile(filepath, entryFunction, "vs_4_0", &m_pBlob)))
			{
				CASSERT(m_pBlob != NULL, "CompileShaderFromFile failed.");
			}

			hr = S_OK;
			if(FAILED(hr = m_pD3DDevice11->CreateVertexShader(m_pBlob->GetBufferPointer(), m_pBlob->GetBufferSize(), NULL, &m_Shader.m_pVertexShader11)))
			{
				CASSERT(m_Shader.m_pVertexShader11 != NULL, "CreateVertexShader failed.");
			}
		}
		break;
	case PIXEL11:
		{
			HRESULT hr = S_OK;
			if(FAILED(hr = CompileShaderFromFile(filepath, entryFunction, "ps_4_0", &m_pBlob)))
			{
				CASSERT(m_pBlob != NULL, "CompileShaderFromFile failed.");
			}

			hr = S_OK;
			if(FAILED(hr = m_pD3DDevice11->CreatePixelShader(m_pBlob->GetBufferPointer(), m_pBlob->GetBufferSize(), NULL, &m_Shader.m_pPixelShader11)))
			{
				CASSERT(m_Shader.m_pPixelShader11 != NULL, "CreateVertexShader failed.");
			}
		}
		break;
	case VERTEX9:
		// TODO - Write this
		break;
	case PIXEL9:
		// TODO - Write this
		break;
	};
}
void Shader::ConstructPixelShader9(LPCWSTR filepath)
{
	m_Type = PIXEL9;

	LPD3DXBUFFER pCode = NULL;
	LPD3DXBUFFER pErrorMsgs = NULL;
	HRESULT hr = S_OK;

	hr = D3DXAssembleShaderFromFile(filepath, NULL, NULL, 0, &pCode, &pErrorMsgs);

	if ((FAILED(hr)) && (pErrorMsgs != NULL))
	{
		char* message = (char*)pErrorMsgs->GetBufferPointer();
    
		MessageBoxA(NULL, message, "Construct Pixel Shader Error", MB_OK);
		CASSERT( pCode != NULL, "Assemble Shader from file failed.");
	}
 
	if (FAILED(hr = m_pD3DDevice9->CreatePixelShader((DWORD*)pCode->GetBufferPointer(), 
                                         &m_Shader.m_pPixelShader9)))
	{
		CASSERT(m_Shader.m_pPixelShader9 != NULL, "CreatePixelShader failed.");
	}	
}

ID3D11PixelShader *Shader::getPixelShader11()  const
{
	CASSERT(m_Type == PIXEL11, "Tried to get wrong type of Shader.");
	return m_Shader.m_pPixelShader11;
}

ID3D11VertexShader *Shader::getVertexShader11()  const
{
	CASSERT(m_Type == VERTEX11, "Tried to get wrong type of Shader.");
	return m_Shader.m_pVertexShader11;
}

IDirect3DPixelShader9 *Shader::getPixelShader9() const
{
	CASSERT(m_Type == PIXEL9, "Tried to get wrong type of Shader.");
	return m_Shader.m_pPixelShader9;
}

IDirect3DVertexShader9 *Shader::getVertexShader9() const
{
	CASSERT(m_Type == VERTEX9, "Tried to get wrong type of Shader.");
	return m_Shader.m_pVertexShader9;
}